import torch
import torch.nn as nn
import torchvision.models as models
from mobilevitv2 import MobileViT  # 假设此文件中定义了MobileViT类

class TripleCombinedModel(nn.Module):
    def __init__(self, num_classes=5, fusion_method='concat'):
        super(TripleCombinedModel, self).__init__()
        self.fusion_method = fusion_method
        
        # 加载预训练的 ResNet-18 模型
        self.resnet18 = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        self.resnet18_features = nn.Sequential(*list(self.resnet18.children())[:-1])
        for param in self.resnet18_features.parameters():
            param.requires_grad = False

        # 加载 MobileViT v2 模型
        self.mobilevit = MobileViT(num_classes=num_classes, variant='mobilevitv2_100')
        self.mobilevit_features = nn.Sequential(*list(self.mobilevit.children())[:-1])
        for param in self.mobilevit_features.parameters():
            param.requires_grad = False
            
        # 加载 MobileNetV3 模型
        self.mobilenetv3 = models.mobilenet_v3_large(weights=models.MobileNet_V3_Large_Weights.IMAGENET1K_V2)
        self.mobilenetv3_features = nn.Sequential(*list(self.mobilenetv3.children())[:-1])
        for param in self.mobilenetv3_features.parameters():
            param.requires_grad = False

        # 动态获取特征维度
        test_input = torch.randn(1, 3, 224, 224)  # 创建测试输入
        with torch.no_grad():
            # 获取 ResNet-18 特征维度
            resnet_out = self.resnet18_features(test_input).view(1, -1)
            num_ftrs_resnet = resnet_out.shape[1]  # 应为 512
            
            # 获取 MobileViT v2 特征维度
            mobilevit_out = self.mobilevit_features(test_input)
            mobilevit_out = nn.AdaptiveAvgPool2d(1)(mobilevit_out).view(1, -1)
            num_ftrs_mobilevit = mobilevit_out.shape[1]  # 应为 128
            
            # 获取 MobileNetV3 特征维度
            mobilenetv3_out = self.mobilenetv3_features(test_input)
            mobilenetv3_out = nn.AdaptiveAvgPool2d(1)(mobilenetv3_out).view(1, -1)
            num_ftrs_mobilenetv3 = mobilenetv3_out.shape[1]  # 应为 960
        
        print(f"[模型维度信息] ResNet特征: {num_ftrs_resnet}, MobileViT特征: {num_ftrs_mobilevit}, MobileNetV3特征: {num_ftrs_mobilenetv3}")
        
        # 根据融合方法确定特征融合维度
        if fusion_method == 'concat':
            combined_dim = num_ftrs_resnet + num_ftrs_mobilevit + num_ftrs_mobilenetv3
            # 分类头输入维度：concat时为3个特征变换后的维度之和
            fc_input_dim = combined_dim * 3
            print(f"[融合维度信息] 特征拼接维度: {combined_dim}，分类头输入维度: {fc_input_dim}")
        elif fusion_method in ['attention', 'weighted_sum']:
            combined_dim = max(num_ftrs_resnet, num_ftrs_mobilevit, num_ftrs_mobilenetv3)
            fc_input_dim = combined_dim
            print(f"[融合维度信息] 特征融合维度: {combined_dim}，分类头输入维度: {fc_input_dim}")
        else:
            raise ValueError(f"不支持的融合方法: {fusion_method}")
        
        # 定义特征变换层（将各模型特征变换到统一维度）
        self.resnet_transform = nn.Linear(num_ftrs_resnet, combined_dim)
        self.mobilevit_transform = nn.Linear(num_ftrs_mobilevit, combined_dim)
        self.mobilenetv3_transform = nn.Linear(num_ftrs_mobilenetv3, combined_dim)
        
        # 定义注意力权重（如果使用注意力融合）
        if fusion_method == 'attention':
            self.attention = nn.Sequential(
                nn.Linear(combined_dim * 3, combined_dim),
                nn.Tanh(),
                nn.Linear(combined_dim, 3),
                nn.Softmax(dim=1)
            )
        
        # 定义分类头（使用动态计算的输入维度）
        self.fc = nn.Sequential(
            nn.Linear(fc_input_dim, 512),
            nn.BatchNorm1d(512),  # 添加BatchNorm提高稳定性
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),  # 随机丢弃50%神经元防止过拟合
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        # 提取各模型特征并调整维度
        resnet_out = self.resnet18_features(x).view(x.size(0), -1)
        mobilevit_out = self.mobilevit_features(x)
        mobilevit_out = nn.AdaptiveAvgPool2d(1)(mobilevit_out).view(x.size(0), -1)
        mobilenetv3_out = self.mobilenetv3_features(x)
        mobilenetv3_out = nn.AdaptiveAvgPool2d(1)(mobilenetv3_out).view(x.size(0), -1)
        
        # 特征变换（映射到统一维度）
        resnet_out = self.resnet_transform(resnet_out)
        mobilevit_out = self.mobilevit_transform(mobilevit_out)
        mobilenetv3_out = self.mobilenetv3_transform(mobilenetv3_out)
        
        # 特征融合逻辑
        if self.fusion_method == 'concat':
            # 直接拼接三个模型的特征
            combined_features = torch.cat([resnet_out, mobilevit_out, mobilenetv3_out], dim=1)
        elif self.fusion_method == 'weighted_sum':
            # 等权重求和融合
            combined_features = (resnet_out + mobilevit_out + mobilenetv3_out) / 3
        elif self.fusion_method == 'attention':
            # 注意力机制融合
            concat_features = torch.cat([resnet_out, mobilevit_out, mobilenetv3_out], dim=1)
            attention_weights = self.attention(concat_features).unsqueeze(2)
            feature_list = torch.stack([resnet_out, mobilevit_out, mobilenetv3_out], dim=2)
            combined_features = torch.bmm(feature_list, attention_weights).squeeze(2)
        
        # 通过分类头进行预测
        output = self.fc(combined_features)
        return output
    
def create_combined_model(num_classes=10, fusion_method='concat'):
    """创建组合模型的工厂函数"""
    return TripleCombinedModel(num_classes=num_classes, fusion_method=fusion_method)

# 模型测试代码
if __name__ == "__main__":
    # 设置测试参数
    num_classes = 102
    test_fusion_method = 'concat'  # 可修改为'attention'或'weighted_sum'
    
    # 创建模型
    model = create_combined_model(
        num_classes=num_classes,
        fusion_method=test_fusion_method
    )
    
    # 生成测试输入
    input_tensor = torch.randn(1, 3, 224, 224)
    
    # 前向传播
    output = model(input_tensor)
    
    # 打印测试结果
    print("\n[模型测试结果]")
    print(f"输入形状: {input_tensor.shape}")
    print(f"输出形状: {output.shape}")
    print(f"分类头第一层输入维度: {model.fc[0].in_features}")
    print(f"分类头最后一层输出维度: {model.fc[-1].out_features}")